package com.example.streamer.ws;

import com.example.streamer.util.ScreenCapture;
import com.example.streamer.util.InputInjector;
import com.example.streamer.api.WindowController;
import com.sun.jna.platform.win32.WinDef.HWND;
import com.sun.jna.Pointer;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.*;
import org.springframework.web.socket.handler.BinaryWebSocketHandler;

import java.awt.Rectangle;
import java.util.concurrent.CopyOnWriteArraySet;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;

@Component
public class ScreenWebSocket extends BinaryWebSocketHandler implements WebSocketHandler {
    private static final CopyOnWriteArraySet<WebSocketSession> sessions = new CopyOnWriteArraySet<>();
    private static final int FPS = 60;
    private static final ObjectMapper objectMapper = new ObjectMapper();
    private volatile int dynamicFps = FPS;
    private volatile float dynamicQuality = 0.7f;

    public ScreenWebSocket() {
        Thread thread = new Thread(() -> {
            while (true) {
                try {
                    if (!sessions.isEmpty()) {
                        System.out.println("准备推送图片，连接数: " + sessions.size());
                        String hwndStr = WindowController.getCurrentHwnd();
                        byte[] img = null;
                        try {
                            if (hwndStr != null) {
                                System.out.println("即将抓取窗口: " + hwndStr);
                                String hwndHex = hwndStr.replace("native@", "");
                                if (hwndHex.startsWith("0x") || hwndHex.startsWith("0X")) {
                                    hwndHex = hwndHex.substring(2);
                                }
                                long hwndVal = Long.parseUnsignedLong(hwndHex, 16);
                                HWND hwnd = new HWND(Pointer.createConstant(hwndVal));
                                img = com.example.streamer.util.ScreenCapture.captureWindowAsJpeg(hwnd, dynamicQuality);
                            } else {
                                System.out.println("未指定窗口，抓取全屏");
                                img = com.example.streamer.util.ScreenCapture.captureScreenAsJpeg(dynamicQuality);
                            }
                        } catch (Exception e) {
                            System.out.println("抓屏异常，跳过本帧");
                            e.printStackTrace();
                            // 避免死循环刷屏
                            Thread.sleep(100);
                            continue;
                        }
                        if (img != null) {
                            System.out.println("推送图片，大小: " + img.length);
                            for (WebSocketSession session : sessions) {
                                try {
                                    session.sendMessage(new BinaryMessage(img));
                                } catch (Exception e) {
                                    e.printStackTrace();
                                }
                            }
                        }
                    }
                    int interval = 1000 / Math.max(1, dynamicFps);
                    Thread.sleep(interval);
                } catch (Exception e) {
                    // 捕获所有异常，防止线程崩溃
                    System.out.println("推流主循环异常，自动恢复");
                    e.printStackTrace();
                    try { Thread.sleep(500); } catch (InterruptedException ignored) {}
                }
            }
        });
        thread.setDaemon(true);
        thread.start();
    }

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        sessions.add(session);
        System.out.println("有客户端连接，当前连接数: " + sessions.size());
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        sessions.remove(session);
    }

    @Override
    public void handleTextMessage(WebSocketSession session, TextMessage message) {
        try {
            JsonNode node = objectMapper.readTree(message.getPayload());
            String type = node.get("type").asText("");
            if ("config".equals(type)) {
                if (node.has("fps")) {
                    dynamicFps = node.get("fps").asInt(FPS);
                }
                if (node.has("quality")) {
                    dynamicQuality = (float) node.get("quality").asDouble(0.7);
                }
                System.out.println("收到动态配置: fps=" + dynamicFps + ", quality=" + dynamicQuality);
                return;
            }
            if ("mouse".equals(type)) {
                String action = node.get("action").asText("");
                int x = node.has("x") ? node.get("x").asInt() : 0;
                int y = node.has("y") ? node.get("y").asInt() : 0;
                int button = 0;
                if (node.has("button")) {
                    int btn = node.get("button").asInt();
                    switch (btn) {
                        case 1: button = java.awt.event.InputEvent.BUTTON1_DOWN_MASK; break;
                        case 2: button = java.awt.event.InputEvent.BUTTON3_DOWN_MASK; break;
                        case 3: button = java.awt.event.InputEvent.BUTTON2_DOWN_MASK; break;
                        default: button = 0;
                    }
                }
                int amount = node.has("amount") ? node.get("amount").asInt() : 0;
                switch (action) {
                    case "move":
                        InputInjector.moveMouse(x, y); break;
                    case "press":
                        if (button != 0) InputInjector.mousePress(button); break;
                    case "release":
                        if (button != 0) InputInjector.mouseRelease(button); break;
                    case "wheel":
                        InputInjector.mouseWheel(amount); break;
                }
            } else if ("keyboard".equals(type)) {
                String action = node.get("action").asText("");
                int keyCode = node.has("keyCode") ? node.get("keyCode").asInt() : 0;
                switch (action) {
                    case "press":
                        InputInjector.keyPress(keyCode); break;
                    case "release":
                        InputInjector.keyRelease(keyCode); break;
                }
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
    }
} 